/**/
#ifndef YOLOV8_TENSORRT_H
#define YOLOV8_TENSORRT_H

#include <NvInfer.h>
#include <cuda_runtime.h>
#include <opencv2/opencv.hpp>
#include <vector>
#include <memory>
#include <string>

int funtion(cv::Mat& img);
static constexpr float ODDS = 0.98f;
class YOLOv8TensorRT {
public:
    YOLOv8TensorRT(const std::string& enginePath);
    ~YOLOv8TensorRT();
    
    // 分类方法 - 返回概率向量
    std::vector<float> classify(const cv::Mat& image);
    
private:
    // TensorRT相关
    std::unique_ptr<nvinfer1::ICudaEngine> engine;
    std::unique_ptr<nvinfer1::IExecutionContext> context;
    cudaStream_t stream;
    
    // 缓冲区
    void* buffers[2];  // 输入和输出
    float* inputBuffer;
    float* outputBuffer;
    
    // 模型参数 - 根据您的训练参数调整
    static constexpr int INPUT_W = 64;       // 根据训练时的imgsz=64调整
    static constexpr int INPUT_H = 64;
    static constexpr int NUM_CLASSES = 4;     // 4个类别（包括背景）
    
    
    // 辅助函数
    void preprocess(const cv::Mat& image, float* inputBuffer);
    bool loadEngine(const std::string& enginePath);
};

// Logger for TensorRT
class Logger : public nvinfer1::ILogger {
    void log(Severity severity, const char* msg) noexcept override {
        if (severity <= Severity::kWARNING)
            std::cout << msg << std::endl;
    }
};

#endif
